#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
学生社交媒体数据可视化模块
"""

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.figure_factory as ff
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

class DataVisualizer:
    def __init__(self, df):
        self.df = df
        self.figures = {}
        
    def create_usage_distribution(self):
        """创建使用时间分布图"""
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('每日使用时间分布', '成瘾评分分布', '心理健康评分分布', '睡眠时间分布'),
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"secondary_y": False}]]
        )
        
        # 使用时间分布
        fig.add_trace(
            go.Histogram(x=self.df['Avg_Daily_Usage_Hours'], name='使用时间', nbinsx=30),
            row=1, col=1
        )
        
        # 成瘾评分分布
        fig.add_trace(
            go.Histogram(x=self.df['Addicted_Score'], name='成瘾评分', nbinsx=10),
            row=1, col=2
        )
        
        # 心理健康评分分布
        fig.add_trace(
            go.Histogram(x=self.df['Mental_Health_Score'], name='心理健康', nbinsx=10),
            row=2, col=1
        )
        
        # 睡眠时间分布
        fig.add_trace(
            go.Histogram(x=self.df['Sleep_Hours_Per_Night'], name='睡眠时间', nbinsx=20),
            row=2, col=2
        )
        
        fig.update_layout(
            title_text="关键指标分布图",
            showlegend=False,
            height=600
        )
        
        self.figures['usage_distribution'] = fig
        return fig
    
    def create_correlation_heatmap(self):
        """创建相关性热力图"""
        numeric_cols = ['Age', 'Avg_Daily_Usage_Hours', 'Sleep_Hours_Per_Night', 
                       'Mental_Health_Score', 'Conflicts_Over_Social_Media', 'Addicted_Score']
        
        corr_matrix = self.df[numeric_cols].corr()
        
        fig = go.Figure(data=go.Heatmap(
            z=corr_matrix.values,
            x=corr_matrix.columns,
            y=corr_matrix.columns,
            colorscale='RdBu',
            zmid=0,
            text=np.round(corr_matrix.values, 2),
            texttemplate="%{text}",
            textfont={"size": 10},
            hoverongaps=False
        ))
        
        fig.update_layout(
            title='变量相关性热力图',
            xaxis_title='变量',
            yaxis_title='变量',
            height=500,
            width=600
        )
        
        self.figures['correlation_heatmap'] = fig
        return fig
    
    def create_platform_analysis(self):
        """创建平台分析图"""
        # 平台使用分布
        platform_counts = self.df['Most_Used_Platform'].value_counts()
        
        fig1 = go.Figure(data=[go.Pie(
            labels=platform_counts.index,
            values=platform_counts.values,
            hole=0.3
        )])
        
        fig1.update_layout(title="社交媒体平台使用分布")
        self.figures['platform_distribution'] = fig1
        
        # 平台与成瘾程度的关系
        platform_addiction = self.df.groupby('Most_Used_Platform')['Addicted_Score'].mean().sort_values(ascending=True)
        
        fig2 = go.Figure(data=[go.Bar(
            x=platform_addiction.values,
            y=platform_addiction.index,
            orientation='h',
            marker_color='skyblue'
        )])
        
        fig2.update_layout(
            title="各平台平均成瘾评分",
            xaxis_title="平均成瘾评分",
            yaxis_title="社交媒体平台",
            height=400
        )
        
        self.figures['platform_addiction'] = fig2
        
        return fig1, fig2
    
    def create_demographic_analysis(self):
        """创建人口统计分析图"""
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('性别分布', '学历分布', '性别-使用时间', '学历-成瘾评分'),
            specs=[[{"type": "domain"}, {"type": "domain"}],
                   [{"type": "xy"}, {"type": "xy"}]]
        )
        
        # 性别分布
        gender_counts = self.df['Gender'].value_counts()
        fig.add_trace(
            go.Pie(labels=gender_counts.index, values=gender_counts.values, name="性别"),
            row=1, col=1
        )
        
        # 学历分布
        academic_counts = self.df['Academic_Level'].value_counts()
        fig.add_trace(
            go.Pie(labels=academic_counts.index, values=academic_counts.values, name="学历"),
            row=1, col=2
        )
        
        # 性别与使用时间
        gender_usage = self.df.groupby('Gender')['Avg_Daily_Usage_Hours'].mean()
        fig.add_trace(
            go.Bar(x=gender_usage.index, y=gender_usage.values, name="使用时间"),
            row=2, col=1
        )
        
        # 学历与成瘾评分
        academic_addiction = self.df.groupby('Academic_Level')['Addicted_Score'].mean()
        fig.add_trace(
            go.Bar(x=academic_addiction.index, y=academic_addiction.values, name="成瘾评分"),
            row=2, col=2
        )
        
        fig.update_layout(
            title_text="人口统计学分析",
            showlegend=False,
            height=600
        )
        
        self.figures['demographic_analysis'] = fig
        return fig
    
    def create_academic_impact_analysis(self):
        """创建学术影响分析图"""
        # 学术表现影响分布
        academic_impact = self.df['Affects_Academic_Performance'].value_counts()
        
        fig1 = go.Figure(data=[go.Bar(
            x=academic_impact.index,
            y=academic_impact.values,
            marker_color=['red' if x == 'Yes' else 'green' for x in academic_impact.index]
        )])
        
        fig1.update_layout(
            title="社交媒体对学术表现的影响",
            xaxis_title="是否影响学术表现",
            yaxis_title="学生数量"
        )
        
        # 使用时间与学术表现的关系
        fig2 = go.Figure()
        
        for status in ['Yes', 'No']:
            data = self.df[self.df['Affects_Academic_Performance'] == status]['Avg_Daily_Usage_Hours']
            fig2.add_trace(go.Box(
                y=data,
                name=f"影响学术表现: {status}",
                boxpoints='outliers'
            ))
        
        fig2.update_layout(
            title="使用时间与学术表现影响的关系",
            yaxis_title="每日使用时间(小时)"
        )
        
        self.figures['academic_impact'] = fig1
        self.figures['usage_academic_relation'] = fig2
        
        return fig1, fig2
    
    def create_mental_health_analysis(self):
        """创建心理健康分析图"""
        # 使用时间与心理健康的散点图
        fig1 = px.scatter(
            self.df, 
            x='Avg_Daily_Usage_Hours', 
            y='Mental_Health_Score',
            color='Gender',
            size='Addicted_Score',
            hover_data=['Age', 'Sleep_Hours_Per_Night'],
            title="使用时间与心理健康关系",
            labels={
                'Avg_Daily_Usage_Hours': '每日使用时间(小时)',
                'Mental_Health_Score': '心理健康评分'
            }
        )
        
        # 添加趋势线
        z = np.polyfit(self.df['Avg_Daily_Usage_Hours'], self.df['Mental_Health_Score'], 1)
        p = np.poly1d(z)
        fig1.add_trace(go.Scatter(
            x=self.df['Avg_Daily_Usage_Hours'],
            y=p(self.df['Avg_Daily_Usage_Hours']),
            mode='lines',
            name='趋势线',
            line=dict(color='red', width=2, dash='dash')
        ))
        
        # 睡眠时间与心理健康的关系
        fig2 = px.scatter(
            self.df,
            x='Sleep_Hours_Per_Night',
            y='Mental_Health_Score',
            color='Addicted_Score',
            title="睡眠时间与心理健康关系",
            labels={
                'Sleep_Hours_Per_Night': '每晚睡眠时间(小时)',
                'Mental_Health_Score': '心理健康评分'
            }
        )
        
        self.figures['mental_health_usage'] = fig1
        self.figures['mental_health_sleep'] = fig2
        
        return fig1, fig2
    
    def create_relationship_analysis(self):
        """创建人际关系分析图"""
        # 恋爱状态分布
        relationship_counts = self.df['Relationship_Status'].value_counts()
        
        fig1 = go.Figure(data=[go.Pie(
            labels=relationship_counts.index,
            values=relationship_counts.values,
            hole=0.3
        )])
        
        fig1.update_layout(title="恋爱关系状态分布")
        
        # 恋爱状态与冲突次数的关系
        fig2 = go.Figure()
        
        for status in self.df['Relationship_Status'].unique():
            data = self.df[self.df['Relationship_Status'] == status]['Conflicts_Over_Social_Media']
            fig2.add_trace(go.Box(
                y=data,
                name=status,
                boxpoints='outliers'
            ))
        
        fig2.update_layout(
            title="恋爱状态与社交媒体冲突次数关系",
            yaxis_title="冲突次数"
        )
        
        self.figures['relationship_distribution'] = fig1
        self.figures['relationship_conflicts'] = fig2
        
        return fig1, fig2
    
    def create_comprehensive_dashboard(self):
        """创建综合仪表板"""
        # 这将创建一个包含多个子图的综合仪表板
        fig = make_subplots(
            rows=3, cols=3,
            subplot_titles=(
                '使用时间分布', '成瘾评分分布', '心理健康评分',
                '平台使用分布', '性别分布', '学历分布',
                '使用时间vs心理健康', '平台成瘾评分', '恋爱状态冲突'
            ),
            specs=[
                [{"type": "histogram"}, {"type": "histogram"}, {"type": "histogram"}],
                [{"type": "domain"}, {"type": "domain"}, {"type": "domain"}],
                [{"type": "scatter"}, {"type": "bar"}, {"type": "box"}]
            ]
        )
        
        # 第一行 - 分布图
        fig.add_trace(
            go.Histogram(x=self.df['Avg_Daily_Usage_Hours'], name='使用时间'),
            row=1, col=1
        )
        
        fig.add_trace(
            go.Histogram(x=self.df['Addicted_Score'], name='成瘾评分'),
            row=1, col=2
        )
        
        fig.add_trace(
            go.Histogram(x=self.df['Mental_Health_Score'], name='心理健康'),
            row=1, col=3
        )
        
        # 第二行 - 饼图
        platform_counts = self.df['Most_Used_Platform'].value_counts()
        fig.add_trace(
            go.Pie(labels=platform_counts.index, values=platform_counts.values, name="平台"),
            row=2, col=1
        )
        
        gender_counts = self.df['Gender'].value_counts()
        fig.add_trace(
            go.Pie(labels=gender_counts.index, values=gender_counts.values, name="性别"),
            row=2, col=2
        )
        
        academic_counts = self.df['Academic_Level'].value_counts()
        fig.add_trace(
            go.Pie(labels=academic_counts.index, values=academic_counts.values, name="学历"),
            row=2, col=3
        )
        
        # 第三行 - 关系图
        fig.add_trace(
            go.Scatter(
                x=self.df['Avg_Daily_Usage_Hours'],
                y=self.df['Mental_Health_Score'],
                mode='markers',
                name='使用时间vs心理健康'
            ),
            row=3, col=1
        )
        
        platform_addiction = self.df.groupby('Most_Used_Platform')['Addicted_Score'].mean()
        fig.add_trace(
            go.Bar(x=platform_addiction.index, y=platform_addiction.values, name="平台成瘾"),
            row=3, col=2
        )
        
        for status in self.df['Relationship_Status'].unique():
            data = self.df[self.df['Relationship_Status'] == status]['Conflicts_Over_Social_Media']
            fig.add_trace(
                go.Box(y=data, name=status),
                row=3, col=3
            )
        
        fig.update_layout(
            title_text="学生社交媒体使用综合分析仪表板",
            showlegend=False,
            height=1200,
            width=1400
        )
        
        self.figures['comprehensive_dashboard'] = fig
        return fig
    
    def save_all_figures(self, output_dir='visualizations'):
        """保存所有图表"""
        import os
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        for name, fig in self.figures.items():
            fig.write_html(f"{output_dir}/{name}.html")
            print(f"保存图表: {name}.html")
    
    def create_all_visualizations(self):
        """创建所有可视化图表"""
        print("正在创建可视化图表...")
        
        self.create_usage_distribution()
        self.create_correlation_heatmap()
        self.create_platform_analysis()
        self.create_demographic_analysis()
        self.create_academic_impact_analysis()
        self.create_mental_health_analysis()
        self.create_relationship_analysis()
        self.create_comprehensive_dashboard()
        
        print(f"完成！共创建了 {len(self.figures)} 个图表")
        return self.figures

# 使用示例
if __name__ == "__main__":
    # 加载数据
    df = pd.read_csv('学生社交媒体与人际关系数据集/学生社交媒体与人际关系数据集.csv')
    
    # 创建可视化器
    visualizer = DataVisualizer(df)
    
    # 创建所有图表
    figures = visualizer.create_all_visualizations()
    
    # 保存图表
    visualizer.save_all_figures()
